package leecode.topk;

import java.util.Arrays;
import java.util.Comparator;
import java.util.HashMap;
import java.util.Map;
import java.util.PriorityQueue;

/**
 * @author wangxi created on 2021/5/16 12:21
 * @version v1.0
 *
 * https://leetcode-cn.com/problems/top-k-frequent-elements/
 *
 * 给你一个整数数组 nums 和一个整数 k ，请你返回其中出现频率前 k 高的元素。你可以按 任意顺序 返回答案。
 */
public class TopKFrequent {
    public static void main(String[] args) {
        TopKFrequent obj = new TopKFrequent();
        int[] nums = {1,1,1,2,2,3};
        int[] res = obj.topKFrequent(nums, 2);

        System.out.println(Arrays.toString(res));
    }

    public int[] topKFrequent(int[] nums, int k) {
        if (nums == null || nums.length <= 0) {
            return new int[0];
        }
        Map<Integer, Integer> dict = new HashMap<>();
        for (int num : nums) {
            dict.put(num, dict.get(num) == null ? 1 : dict.get(num) + 1);
        }
        // 建立最小堆，因此按照小到大排序，堆顶元素即为最小元素，每次poll会去掉最小元素
        PriorityQueue<Integer> queue = new PriorityQueue<>(Comparator.comparingInt(dict::get));
        for (int item : dict.keySet()) {
            queue.offer(item);
            if (queue.size() > k) {
                // 自动弹出最小元素
                queue.poll();
            }
        }
        int[] res = new int[queue.size()];
        int i = 0;
        // 堆可以这样遍历，这种是小到大排序
        while (!queue.isEmpty()) {
            res[i] = queue.poll();
            i++;
        }
        return res;
    }
}
